from __future__ import division, print_function, absolute_import

import tensorflow.compat.v1 as tf

from cnn_estimator import cnn_model_fn
from cifar10 import input_fn

if __name__ == '__main__':
    tf.logging.set_verbosity(tf.logging.INFO)
    params = {
        'learning_rate': 0.001,
        'dropout_rate': 0.4,
        'print_shape': True
    }
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5)
    config = tf.ConfigProto(gpu_options=gpu_options)
    run_config = tf.estimator.RunConfig(
        model_dir='./cnn_estimator/',
        save_summary_steps=20,
        save_checkpoints_steps=500,
        session_config=config)
    estimator = tf.estimator.Estimator(model_fn=cnn_model_fn,
                                       config=run_config,
                                       params=params)

    train_file = '../data/train.tfrecord'
    eval_file = '../data/test.tfrecord'
    train_batch_size = 256
    eval_batch_size = 64
    train_steps = 40000 // train_batch_size

    train_input_fn = input_fn(train_file, train_batch_size)
    eval_input_fn = input_fn(eval_file, eval_batch_size)

    # model.train(input_fn=train_input_fn, steps=2000)

    train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=10000)
    eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
